package com.hk.core.web.reactive.method.annotations;

import com.hk.commons.annotations.QueryModelRequestBody;
import com.hk.commons.query.OrderUtils;
import com.hk.commons.query.QueryModel;
import org.springframework.core.MethodParameter;
import org.springframework.core.ReactiveAdapterRegistry;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.util.Assert;
import org.springframework.web.reactive.BindingContext;
import org.springframework.web.reactive.result.method.annotation.AbstractMessageReaderArgumentResolver;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.List;

/**
 * 对QueryModel 中的Order字段进行过滤，防止SQL 注入.
 *
 * @author Kevin
 * @date 2024-03-26 08:53
 */
public class QueryModelRequestBodyMethodArgumentResolver extends AbstractMessageReaderArgumentResolver {

    public QueryModelRequestBodyMethodArgumentResolver(List<HttpMessageReader<?>> readers, ReactiveAdapterRegistry registry) {
        super(readers, registry);
    }

    @Override
    public boolean supportsParameter(MethodParameter parameter) {
        return QueryModel.class.isAssignableFrom(parameter.getParameterType()) && parameter.hasParameterAnnotation(QueryModelRequestBody.class);
    }

    @Override
    public Mono<Object> resolveArgument(MethodParameter parameter, BindingContext bindingContext, ServerWebExchange exchange) {
        QueryModelRequestBody ann = parameter.getParameterAnnotation(QueryModelRequestBody.class);
        Assert.state(ann != null, "No QueryModelRequestBody annotation");
        return readBody(parameter, ann.required(), bindingContext, exchange);
    }

    @Override
    protected Mono<Object> readBody(MethodParameter bodyParam, MethodParameter actualParam, boolean isBodyRequired,
                                    BindingContext bindingContext, ServerWebExchange exchange) {
        return super.readBody(bodyParam, actualParam, isBodyRequired, bindingContext, exchange)
                .map(obj -> {
                    if (obj instanceof QueryModel<?> queryModel) {
                        queryModel.setOrders(OrderUtils.escapeFilter(queryModel.getOrders()));
                        return queryModel;
                    }
                    return obj;
                });
    }
}
